import os
import json
import numpy as np
from .localizer import Localizer


class SuperTransform:
    def __init__(self, local_csv, calib_json):
        assert os.path.exists(local_csv), '{} does not exist!'.format(local_csv)
        assert os.path.exists(calib_json), '{} does not exist!'.format(calib_json)
        self.localizer = Localizer(csv_file=local_csv)
        self.static_transfoms = dict()
        with open(calib_json, 'r') as f:
            json_data = json.load(f)
            for key, val in json_data.items():
                if isinstance(val, dict) and 'transforms' in val.keys():
                    for skey, sval in val['transforms'].items():
                        self.static_transfoms[key] = {skey: np.array(sval)}
        pass

    def get_tf_from_time(self, source_time, target_time):
        source2local = self.localizer.get_tf(stamp=source_time)
        target2local = self.localizer.get_tf(stamp=target_time)
        source2target = np.linalg.inv(target2local) @ source2local
        return source2target

    def get_tf_from_space(self, source_point, target_point='base_link'):
        # 保证每个节点和baselink连接
        assert source_point in self.static_transfoms, f'{source_point} not in static_transfoms!'
        if target_point == 'base_link':
            return self.static_transfoms[source_point]['base_link']
        return (np.linalg.inv(self.get_tf_from_space(source_point=target_point, target_point='base_link')) @
                self.get_tf_from_space(source_point=source_point, target_point='base_link'))

    def get_tf(self, source_time, source_point, target_time, target_point='base_link'):
        past_base2now_base = self.get_tf_from_time(source_time=source_time, target_time=target_time)
        sensor2base = self.get_tf_from_space(source_point=source_point)
        if past_base2now_base is None or sensor2base is None:
            return None
        past_sensor2now_base = past_base2now_base @ sensor2base
        return past_sensor2now_base


if __name__ == '__main__':
    super_transform = SuperTransform(
        local_csv='/media/adt/ZWH4T/ZWH/bags/dataset/sll/test_out/__roate/localization/localization.csv',
        calib_json='/media/adt/ZWH4T/ZWH/bags/dataset/sll/test_out/__roate/calib/calib.json')
    t12t0 = super_transform.get_tf_from_time(source_time=1685296485760, target_time=1685296485760)
    l2b = super_transform.get_tf_from_space(source_point='lidar3')
    l2t0 = super_transform.get_tf(source_time=1685296485760, source_point='lidar3', target_time=1685296487960)
    pass
